/*
卷积神经网络层头文件
定义卷积操作、特征提取和参数管理
支持多种卷积核大小和步长配置
*/
#ifndef CONVLAYER_H
#define CONVLAYER_H

#include "Macros.h"
#ifdef USE_USER_DEFINED_TENSOR 
    #undef USE_EIGEN_TENSOR
    #include "Tensor.h"
#endif
#ifdef USE_EIGEN_TENSOR
    #undef USE_USER_DEFINED_TENSOR
    #include <Eigen/Core>
    #include <Eigen/Dense>
#endif
#include "Layer.h"
#include "LayerManager.h"
#include "Filter.h"
#include "ActivationFunction.h"
#include "LossFunction.h"
#include "PoolingLayer.h"
#include <vector>
#include <memory>
#include <string>

#ifdef USE_EIGEN_TENSOR
// Enable Eigen parallelization
#define EIGEN_USE_THREADS
#endif

using namespace std;
#ifdef USE_EIGEN_TENSOR
using namespace Eigen;
#endif
#ifdef USE_USER_DEFINED_TENSOR
using namespace UserDefinedTensor;
#endif

class ConvLayer;
class PoolingLayer;
class FullConnectedLayer;

void zeroPadding(const Tensor<double, 3>& inputFeatureMap, Tensor<double, 3>& outputFeatureMap, int padding);
void conv2d(const Tensor<double, 3>& inputFeatureMap, const Filter& filter, Tensor<double, 3>& outputFeatureMap, int widthStride, int heightStride);
void convolve(const Tensor<double, 3>& inputFeatureMap, const Tensor<double, 3>& filter, Tensor<double, 3>& outputFeatureMap, int widthStride, int heightStride, int padding);

class ConvLayer : public Layer {
public:
    ConvLayer(int layerIndex, const LayerType& layerType, int inputWidth, int inputHeight, int inputChannels, int numFilters, int filterWidth, int filterHeight, int widthStride, int heightStride, int padding, const string& activationName = "relu", const string& lossName = "mse");
    ~ConvLayer();
    void setPrevLayer(const shared_ptr<Layer>& prevLayer) override;
    void setNextLayer(const shared_ptr<Layer>& nextLayer) override;
    const LayerType& getLayerType() const;
    int getInputWidth() const;
    int getInputHeight() const;
    int getInputChannels() const;
    const Tensor<double, 3>& getInputFeatureMap() const;
    int getNumFilters() const;
    const vector<shared_ptr<Filter> >& getFilters() const;
    int getFilterWidth() const;
    int getFilterHeight() const;
    int getWidthStride() const;
    int getHeightStride() const;
    int getPadding() const;
    int getOutputWidth() const;
    int getOutputHeight() const;
    int getOutputChannels() const;
    int getOutputWidthForPooling() const;
    int getOutputHeightForPooling() const;
    int getOutputChannelsForPooling() const;
    int getOutputSizeForFc() const;
    const Tensor<double, 3>& getOutputFeatureMap() const;
    const Tensor<double, 3>& getOutputFeatureMapForPooling() const;
    const Tensor<double, 1>& getOutputFeatureMapForFc() const;
    const Tensor<double, 3> getInputsGradient() const;
    const Tensor<double, 3> getOutputsGradient() const;
    const Tensor<double, 3> getWeightsGradient() const;
    const Tensor<double, 1> getBiasesGradient() const;
    const Tensor<double, 3> getInputsDelta() const;
    const Tensor<double, 3> getOutputsDelta() const;
    const Tensor<double, 3> getWeightsDelta() const;
    const Tensor<double, 1> getBiasesDelta() const;
    void setLayerType(const LayerType& layerType);
    void setInputFeatureMap(const Tensor<double, 3>& inputFeatureMap);
    void setInputWidth(int inputWidth);
    void setInputHeight(int inputHeight);
    void setInputChannels(int inputChannels);
    void setNumFilters(int numFilters);
    void setFilters(const vector<shared_ptr<Filter> >& filters);
    void setFilterWidth(int filterWidth);
    void setFilterHeight(int filterHeight);
    void setWidthStride(int widthStride);
    void setHeightStride(int heightStride);
    void setPadding(int padding);
    void setOutputWidth(int outputWidth);
    void setOutputHeight(int outputHeight);
    void setOutputChannels(int outputChannels);
    void setOutputWidthForPooling(int outputWidthForPooling);
    void setOutputHeightForPooling(int outputHeightForPooling);
    void setOutputChannelsForPooling(int outputChannelsForPooling);
    void setOutputSizeForFc(int outputSizeForFc);
    void calculateOutputDimensions(int inputWidth, int inputHeight, int filterWidth, int filterHeight, int widthStride, int heightStride, int padding);
    void calculateOutputFeatureMap();
    void transformOutputFeatureMapToPooling(Tensor<double, 3>& outputFeatureMap);
    void transformOutputFeatureMapToFc(Tensor<double, 3>& outputFeatureMap);
    void calculateInputsGradient();
    void calculateOutputsGradient(const Tensor<double, 3>& nextLayerOutputsGradient);
    void calculateWeightsGradient();
    void calculateBiasesGradient();
    void calculateInputsDelta(double learningRate, double momentum);
    void calculateOutputsDelta(double learningRate, double momentum);
    void calculateWeightsDelta(double learningRate, double momentum);
    void calculateBiasesDelta(double learningRate, double momentum);
    void forward(const Tensor<double, 3>& inputFeatureMap);
    void backward(const Tensor<double, 3>& nextLayerOutputsGradient);
    void updateWeights(double learningRate, double momentum);
    void updateBiases(double learningRate, double momentum);
    void print() const;

private:
    Tensor<double, 3> inputFeatureMap;
    int inputWidth;
    int inputHeight;
    int inputChannels;
    int numFilters;
    vector<shared_ptr<Filter> > filters;
    int filterWidth;
    int filterHeight;
    int widthStride;
    int heightStride;
    int padding;
    Tensor<double, 3> outputFeatureMap; // 3D tensor for Layer outputs
    Tensor<double, 3> outputFeatureMapForPooling;
    Tensor<double, 1> outputFeatureMapForFc;
    int outputWidth;
    int outputHeight;
    int outputChannels;
    int outputWidthForPooling;
    int outputHeightForPooling;
    int outputChannelsForPooling;
    int outputSizeForFc;
    Tensor<double, 3> inputsGradient;
    Tensor<double, 3> outputsGradient;
    Tensor<double, 3> weightsGradient;
    Tensor<double, 1> biasesGradient;
    Tensor<double, 3> inputsDelta;
    Tensor<double, 3> outputsDelta;
    Tensor<double, 3> weightsDelta;
    Tensor<double, 1> biasesDelta;
    Tensor<double, 3> prevBatchInputsDelta;
    Tensor<double, 3> prevBatchOutputsDelta;
    Tensor<double, 3> prevBatchWeightsDelta;
    Tensor<double, 1> prevBatchBiasesDelta;
    unique_ptr<ActivationFunction<double, 3> > activation;
    unique_ptr<LossFunction> loss;
    shared_ptr<ConvLayer> prevConvLayer;
    shared_ptr<PoolingLayer> prevPoolingLayer;
    shared_ptr<PoolingLayer> nextPoolingLayer;
    shared_ptr<FullConnectedLayer> nextFcLayer;
};

#endif // CONVLAYER_H